package com.smart.thrift.discoveryandbalance.discovery.zookeeper;

import cn.hutool.json.JSONUtil;
import com.google.common.collect.Sets;
import com.smart.thrift.constants.ThriftConstants;
import com.smart.thrift.discoveryandbalance.discovery.AbstractRegistry;
import com.smart.thrift.discoveryandbalance.discovery.ThriftServerNode;
import com.smart.thrift.exception.client.ThriftClientRegistryException;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.retry.RetryNTimes;
import org.apache.zookeeper.CreateMode;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;
import java.util.Set;

@Slf4j
public class ZookeeperRegistry extends AbstractRegistry<CuratorFramework> {

    public static final String ZOOKEEPER_NODE_PREFIX = "/thrift-node";

    private CuratorFramework zkClient;

    @Getter
    private final String clientType = ThriftConstants.DiscoveryType.ZOOKEEPER;

    private static final Integer SESSION_TIMEOUT = 5000;
    private static final Integer CONNECTION_TIMEOUT = 5000;
    private static final Integer RETRY_TIMES = 10;
    private static final Integer SLEEP_MS_BETWEEN_RETRIES = 5000;

    private static ZookeeperRegistry zookeeperRegistry;

    private ZookeeperRegistry() {}

    public static ZookeeperRegistry getInstance(){
        if (Objects.isNull(zookeeperRegistry)) {
            zookeeperRegistry = new ZookeeperRegistry();
            zookeeperRegistry.createClient();
            zookeeperRegistry.startClient();
        }
        return zookeeperRegistry;
    }

    private String getServicePath(String serviceName, String serverId) {
        return ZOOKEEPER_NODE_PREFIX + "/" + serviceName + "/" + serverId;
    }

    @Override
    public CuratorFramework getClient() {
        return zkClient;
    }

    @Override
    public void register(String serviceName, String serverId, ThriftServerNode thriftServerNode) throws Exception {
        if (zkClient.getState() != CuratorFrameworkState.STARTED) {
            zkClient.start();
            zkClient.getZookeeperClient().blockUntilConnectedOrTimedOut();
        }
        String strNodeInfo = JSONUtil.toJsonStr(thriftServerNode);
        zkClient.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT)
                .forPath(getServicePath(serviceName, serverId), strNodeInfo.getBytes(StandardCharsets.UTF_8));
    }

    @Override
    public ThriftServerNode unregister(String serviceName, String serverId) throws Exception {
        if (zkClient.getState() != CuratorFrameworkState.STARTED) {
            zkClient.start();
            zkClient.getZookeeperClient().blockUntilConnectedOrTimedOut();
        }
        ThriftZookeeperServerNode registeredZkNode = getRegisteredZkNode(serviceName, serverId);
        String nodePath = getServicePath(serviceName, serverId);
        if (null != zkClient.checkExists().forPath(nodePath)) {
            zkClient.delete().forPath(nodePath);
        }
        return registeredZkNode;
    }

    public void startClient(){
        zkClient.start();
    }

    public void createClient() {
        if(Objects.nonNull(zkClient)) {
            closeClient();
        }
        try {
            zkClient = CuratorFrameworkFactory.newClient(getRegistryProperties().getConnectString(), SESSION_TIMEOUT, CONNECTION_TIMEOUT,
                    new RetryNTimes(RETRY_TIMES, SLEEP_MS_BETWEEN_RETRIES));
        } catch (Exception e) {
            throw new ThriftClientRegistryException("Unable to access consul server, address is: " + getRegistryProperties().getConnectString(), e);
        }

        if (Objects.isNull(zkClient)) {
            throw new ThriftClientRegistryException("Unable to access consul server, address is: " + getRegistryProperties().getConnectString());
        }
    }

    public void closeClient() {
        CuratorFramework client = zkClient;
        zkClient = null;
        if (client != null) {
            try {
                client.close();
            } catch (Exception e) {
                e.printStackTrace();
                log.warn("Exception when closing zookeeper client, " + e);
            }
        }
    }

    public ThriftZookeeperServerNode getRegisteredZkNode(String serviceName, String serverId) throws Exception {
        String nodePath = getServicePath(serviceName, serverId);
        return JSONUtil.toBean(new String(zkClient.getData().forPath(nodePath)), ThriftZookeeperServerNode.class);
    }

    public Set<ThriftZookeeperServerNode> serviceNodes(List<String> addresses) throws Exception {
        Set<ThriftZookeeperServerNode> addressList = Sets.newHashSetWithExpectedSize(addresses.size());
        for (String node : addresses) {
            String nodeData = new String(zkClient.getData().forPath(ZOOKEEPER_NODE_PREFIX + "/" + node));
            addressList.add(JSONUtil.toBean(nodeData, ThriftZookeeperServerNode.class));
        }
        return addressList;
    }
}
